import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import argparse
# from networks.ClassicNetwork.VGGNet import VGGNet
from networks.ClassicNetwork.ResNet import ResNet50
import os
import numpy as np
import time
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import argparse
# from networks.ClassicNetwork.VGGNet import VGGNet
from networks.ClassicNetwork.ResNet import ResNet50
import os
import numpy as np
import time
from dataloader import load_dataset
from dataloader import load_dataset
from PIL import Image
import torch
from torchvision import transforms
import matplotlib.pyplot as plt
import matplotlib.pyplot as plt
plt.ion()  # interactive mode

# 模型存储路径 只改这两个就行
data_dir = "ResNet50"
#

# ------------------------ 加载数据 --------------------------- #
# Data augmentation and normalization for training
# Just normalization for validation
# 定义预训练变换
preprocess_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

class_names = [ '1', '2','3','4','5','6','7','8','9','10']  # 这个顺序很重要，要和训练时候的类名顺序一致
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
EPOCH = 63  # 最大epoch数目
pre_epoch = 0  # 已训练epoch次数
BATCH_SIZE = 21  # Batchsize
LR = 2e-4  # 学习率
WEIGHT_DECAY = 5e-4  # 衰减系数
STEP_SIZE = 50  # 学习率衰减过程
GAMMA = 0.1  # The decay multiple in each decay step

# ------------------------ 载入模型并且训练 --------------------------- #
net = ResNet50(num_classes=11).to(device)
criterion = nn.CrossEntropyLoss()  # 定义损失函数
optimizer = optim.Adam(net.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=STEP_SIZE, gamma=GAMMA)

checkpoint = torch.load('./model/' + data_dir + '/net.pkl')
net.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
pre_epoch = checkpoint['EPOCH'] - 1
print("已训练epoch：", pre_epoch, "总epoch：", EPOCH)
net.eval()
print("Start Continue Training ResNet50")

# model = torch.load(model_save_path)
# model.eval()
# print(model)
import os
path1 = R'test'
path_list = os.listdir(path1)
# path_list.remove('.DS_Store')  # macos中的文件管理文件，默认隐藏，这里可以忽略
# print(path_list)
#
for i in (path_list):
    path = os.path.join(path1,i)
    
    image_PIL = Image.open(path)
    image_tensor = preprocess_transform(image_PIL)
    # 以下语句等效于 image_tensor = torch.unsqueeze(image_tensor, 0)
    image_tensor.unsqueeze_(0)
    # 没有这句话会报错
    image_tensor = image_tensor.to(device)
    
#     apple_pie': 0, 'baby_back_ribs': 1, 'baklava': 2, 'beef_carpaccio': 3, 'beef_tartare': 4, 'beet_salad': 5, 'beignets': 6, 'bibimbap': 7, 'bread_pudding': 8, 'breakfast_burrito': 9, 'bruschetta': 10
    
    """根据已载入的模型进行识别"""
    class_dict = {
        "0": "apple_pie",
        "1": "baby_back_ribs",
        "2": "baklava",
        "3": "beef_carpaccio",
        "4": "beef_tartare",
        "5": "beet_salad",
        "6": "beignets",
        "7": "bibimbap",
        "8": "bread_pudding",
        "9": "breakfast_burrito",
        "10": "bruschetta"

    }
    out, features = net(image_tensor)
    # print("out", out)
    # 得到预测结果，并且从大到小排序
    preValue, predicted = torch.max(out.data, 1)

    predicted = predicted.cpu().numpy()
    if predicted == 0:
        continue
    class_str = class_dict[str(predicted[0])]
#     print(class_str)
#     print(path)
    print(path,"预测类别为第prediction:{}".format(class_str))
    # print("preValue:{}".format(preValue))
#     _, indices = torch.sort(out, descending=True)
    # 返回每个预测值的百分数
#     percentage = torch.nn.functional.softmax(out, dim=1)[0] 
# 	print（"预测概率",percentage）
# 	plt.imshow(image_tensor)
# 	plt.axis("off")
# 	text = 'percentage: {:.2f}'.format(percentage)
#     if text:
#         plt.text(75, 8, text, style='italic',fontweight='bold',
#             bbox={'facecolor':'white', 'alpha':0.8, 'pad':10})
# 	plt.savefig("1.jpg")
#     plt.show()  
	
	
	
	
	

    # print([(class_names[idx], percentage[idx].item()) for idx in indices[0][:5]])